# -- coding: utf-8 --
import msvcrt
import os
import numpy as np
import cv2
from Mv3dRgbdImport.Mv3dRgbdDefine import *
from Mv3dRgbdImport.Mv3dRgbdApi import *
from Mv3dRgbdImport.Mv3dRgbdDefine import DeviceType_Ethernet, DeviceType_USB, DeviceType_Ethernet_Vir, DeviceType_USB_Vir
from ctypes import *

# 全局变量用于保存图像数据
global_color_image = None
global_depth_image = None

FrameInfoCallBack = WINFUNCTYPE(None, c_void_p, c_void_p)

def get_depth_at_point(depth_image, x, y):
    """获取指定坐标的深度值"""
    if depth_image is None:
        print("深度图像未初始化")
        return None
    
    if x < 0 or x >= depth_image.shape[1] or y < 0 or y >= depth_image.shape[0]:
        print("坐标超出图像范围")
        return None
    
    depth_value = depth_image[y, x]
    print(f"坐标({x}, {y})处的深度值为: {depth_value} mm")
    return depth_value

def save_images(color_img, depth_img, prefix="frame"):
    """保存彩色和深度图像"""
    if color_img is not None:
        color_path = f"{prefix}_color.png"
        cv2.imwrite(color_path, color_img)
        print(f"已保存彩色图像到: {color_path}")
    
    if depth_img is not None:
        # 归一化深度图像以便保存
        depth_normalized = cv2.normalize(depth_img, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)
        depth_path = f"{prefix}_depth.png"
        cv2.imwrite(depth_path, depth_normalized)
        print(f"已保存深度图像到: {depth_path}")

def yuv2rgb(np_color,height,width):
   # 方法1：直接转换（推荐）
    try:
        yuyv_image = np_color.reshape((height, width * 2))
        rgb_image = cv2.cvtColor(yuyv_image, cv2.COLOR_YUV2BGR_YUY2)
    except Exception as e:
        # print(f"直接转换失败: {e}")
        
        # 方法2：精确手动转换
        # YUYV格式布局：Y0 U0 Y1 V0 Y2 U2 Y3 V2...
        # 总数据量：width*height*2 bytes
        
        # 创建完整数组
        yuyv_array = np_color.reshape(-1, 4)  # 每组4字节: Y0 U0 Y1 V0
        
        # 提取Y分量（所有像素）
        y = np.zeros(height * width, dtype=np.uint8)
        y[0::2] = yuyv_array[:, 0]  # 偶数像素Y
        y[1::2] = yuyv_array[:, 2]  # 奇数像素Y
        
        # 提取U和V分量（每两个像素共享一组UV）
        u = yuyv_array[:, 1]
        v = yuyv_array[:, 3]
        
        # 调整UV分量尺寸以匹配Y分量
        u = np.repeat(u, 2)  # 将U值重复给两个像素
        v = np.repeat(v, 2)  # 将V值重复给两个像素
        
        # 重组为YUV格式（需要转置为OpenCV期望的格式）
        y = y.reshape(height, width)
        u = u.reshape(height, width)
        v = v.reshape(height, width)
        
        # 合并为3通道图像
        yuv_image = np.dstack((y, u, v)).astype(np.uint8)
        
        # 转换为BGR
        rgb_image = cv2.cvtColor(yuv_image, cv2.COLOR_YUV2BGR)
    
    return rgb_image

def image_callback(pstFrameData, pUser):
    global global_color_image, global_depth_image
    
    if pstFrameData == None:
        print("no data")
        return
    
    stFrameData = cast(pstFrameData, POINTER(MV3D_RGBD_FRAME_DATA)).contents
    
    if stFrameData.nImageCount != 0:
        color_img = None
        depth_img = None
        
        for i in range(0, stFrameData.nImageCount):
            print(f"MV3D_RGBD_FetchFrame[{i}]: nFrameNum[{stFrameData.stImageData[i].nFrameNum}], "
                  f"nDataLen[{stFrameData.stImageData[i].nDataLen}], "
                  f"nWidth[{stFrameData.stImageData[i].nWidth}], "
                  f"nHeight[{stFrameData.stImageData[i].nHeight}]")
            
            print('矫正：', stFrameData.stImageData[i].bIsRectified)
            print('图像格式：', stFrameData.stImageData[i].enImageType)
            
            # 处理深度图像
            if stFrameData.stImageData[i].enImageType == 17825976:
                depth_data = stFrameData.stImageData[i].pData
                height = stFrameData.stImageData[i].nHeight
                width = stFrameData.stImageData[i].nWidth
                
                # 将深度数据转换为numpy数组
                np_depth = np.frombuffer((c_ubyte * stFrameData.stImageData[i].nDataLen).from_address(
                    ctypes.addressof(depth_data.contents)), dtype=np.uint16)
                np_depth = np_depth.reshape((height, width))
                
                # 保存深度图像
                depth_img = np_depth.copy()
                global_depth_image = depth_img
                
                # 显示深度图像
                depth_vis = cv2.normalize(np_depth, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)
                cv2.imshow('Depth Image', depth_vis)
            
            # 处理彩色图像 (YUV422格式)
            elif stFrameData.stImageData[i].enImageType == 34603058:
                color_data = stFrameData.stImageData[i].pData
                height = stFrameData.stImageData[i].nHeight
                width = stFrameData.stImageData[i].nWidth
                
                # 将YUV422数据转换为numpy数组
                np_color = np.frombuffer((c_ubyte * stFrameData.stImageData[i].nDataLen).from_address(
                    ctypes.addressof(color_data.contents)), dtype=np.uint8)
                
                # # YUV422转RGB
                # yuv_image = np_color.reshape((height, width * 2))
                # # rgb_image = cv2.cvtColor(yuv_image, cv2.COLOR_YUV2BGR_YUYV)
                # yuv420_image = cv2.resize(yuv_image, (width, height*2))
                # rgb_image = cv2.cvtColor(yuv420_image, cv2.COLOR_YUV2BGR_I420)
                rgb_image=yuv2rgb(np_color,height,width)
                
                # 保存彩色图像
                color_img = rgb_image.copy()
                global_color_image = color_img
                
                # 显示彩色图像
                cv2.imshow('Color Image', rgb_image)
        
        # 保存图像
        if color_img is not None and depth_img is not None:
            save_images(color_img, depth_img)
        
        key = cv2.waitKey(1)
        if key == ord('q'):
            return
        elif key == ord('d'):  # 按'd'键获取深度信息
            x = int(input("请输入X坐标: "))
            y = int(input("请输入Y坐标: "))
            get_depth_at_point(global_depth_image, x, y)
            z = input("输入任意继续: ")

CALL_BACK_FUN = FrameInfoCallBack(image_callback)

if __name__ == "__main__":
    nDeviceNum = ctypes.c_uint(0)
    nDeviceNum_p = byref(nDeviceNum)
    
    # 获取设备数量
    ret = Mv3dRgbd.MV3D_RGBD_GetDeviceNumber(DeviceType_Ethernet | DeviceType_USB | DeviceType_Ethernet_Vir | DeviceType_USB_Vir, nDeviceNum_p)
    if ret != 0:
        print("MV3D_RGBD_GetDeviceNumber fail! ret[0x%x]" % ret)
        os.system('pause')
        sys.exit()
    if nDeviceNum == 0:
        print("find no device!")
        os.system('pause')
        sys.exit()
    print("Find devices numbers:", nDeviceNum.value)

    # 枚举设备信息
    stDeviceList = MV3D_RGBD_DEVICE_INFO_LIST()
    net = Mv3dRgbd.MV3D_RGBD_GetDeviceList(DeviceType_Ethernet | DeviceType_USB | DeviceType_Ethernet_Vir | DeviceType_USB_Vir, 
                                    pointer(stDeviceList.DeviceInfo[0]), 20, nDeviceNum_p)
    for i in range(0, nDeviceNum.value):
        print("\ndevice: [%d]" % i)
        strModeName = ""
        for per in stDeviceList.DeviceInfo[i].chModelName:
            strModeName = strModeName + chr(per)
        print("device model name: %s" % strModeName)

        strSerialNumber = ""
        for per in stDeviceList.DeviceInfo[i].chSerialNumber:
            strSerialNumber = strSerialNumber + chr(per)
        print("device SerialNumber: %s" % strSerialNumber)

    # 创建相机示例
    camera = Mv3dRgbd()
    nConnectionNum = input("please input the number of the device to connect:")
    if int(nConnectionNum) >= nDeviceNum.value:
        print("intput error!")
        os.system('pause')
        sys.exit()
    
    # 打开设备
    ret = camera.MV3D_RGBD_OpenDevice(pointer(stDeviceList.DeviceInfo[int(nConnectionNum)]))
    if ret != 0:
        print("MV3D_RGBD_OpenDevice fail! ret[0x%x]" % ret)
        os.system('pause')
        sys.exit()

    # 注册回调函数
    camera.MV3D_RGBD_RegisterFrameCallBack(CALL_BACK_FUN, None)

    # 开始取流
    ret = camera.MV3D_RGBD_Start()
    if ret != 0:
        print("start fail! ret[0x%x]" % ret)
        camera.MV3D_RGBD_CloseDevice()
        os.system('pause')
        sys.exit()
        
    print("按 'q' 键退出程序")
    print("按 'd' 键获取指定坐标的深度信息")
    
    # 等待用户退出
    while True:
        if msvcrt.kbhit() and msvcrt.getch() == b'q':
            break
    
    # 停止取流
    ret = camera.MV3D_RGBD_Stop()
    if ret != 0:
        print("stop fail! ret[0x%x]" % ret)
        os.system('pause')
        sys.exit()

    # 销毁句柄
    ret = camera.MV3D_RGBD_CloseDevice()
    if ret != 0:
        print("CloseDevice fail! ret[0x%x]" % ret)
        os.system('pause')
        sys.exit()
        
    cv2.destroyAllWindows()
    sys.exit()